# -*- coding: UTF-8 -*-
"""
@Project , trainee
@File    , type_year_analysisData.py
@IDE     , PyCharm
@Author  , 2607750505@qq.com
@Date    , 2025/6/17 15:56
"""
import os
import pandas as pd
from tools.getDataBase import get_conn


def get_movie_types():
    """获取所有电影类型（用于搜索建议）"""
    conn, cursor = get_conn()
    cursor.execute("SELECT DISTINCT types FROM movies")
    results = cursor.fetchall()
    conn.close()

    # 提取并合并所有类型
    all_types = set()
    for row in results:
        if row[0]:  # 检查是否为空
            all_types.update(row[0].split(','))
    return sorted(list(all_types))


def get_type_year_data(target_type=None):
    """获取指定类型的历年数量统计"""
    conn, cursor = get_conn()

    try:
        # 获取所有电影的年份和类型
        cursor.execute("SELECT year, types FROM movies")
        movies = cursor.fetchall()

        # 处理年份和类型数据
        year_counts = {}
        for movie in movies:
            year = str(movie[0])[:4]  # 提取年份前4位
            types = movie[1].split(',') if movie[1] else []

            # 如果没有指定类型，或包含目标类型，则统计
            if not target_type or target_type in types:
                year_counts[year] = year_counts.get(year, 0) + 1

        # 按年份排序
        sorted_years = sorted(year_counts.keys())
        years = sorted_years
        counts = [year_counts[year] for year in sorted_years]

        return years, counts

    finally:
        if conn:
            conn.close()
